"""
修复轨迹跟踪问题的测试脚本
"""
import numpy as np
import matplotlib.pyplot as plt
import os

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from utils import set_plot_style

def test_simple_circle():
    """测试简单的圆形轨迹跟踪"""
    print("=" * 50)
    print("测试简单的圆形轨迹跟踪")
    print("=" * 50)
    
    # 设置参数
    dt = 0.1
    sim_time = 30.0
    time_points = np.arange(0, sim_time, dt)
    
    # 创建轨迹生成器并生成圆形轨迹
    traj_gen = TrajectoryGenerator(dt=dt)
    trajectory = []
    
    # 指定圆的参数
    radius = 5.0
    center = (5.0, 5.0)
    omega = 0.3
    
    print(f"生成圆形轨迹: 半径={radius}, 圆心={center}, 角速度={omega}")
    
    for t in time_points:
        state = traj_gen.circle_trajectory(t, radius=radius, center=center, omega=omega)
        trajectory.append(state)
    
    trajectory = np.array(trajectory)
    
    # 创建车辆模型
    vehicle = DoubleIntegratorVehicle(dt=dt)
    
    # 获取系统矩阵
    A, B = vehicle.get_state_space_matrices()
    
    # 创建LQR控制器
    Q = np.diag([1000.0, 1000.0, 100.0, 100.0])
    R = np.diag([0.01, 0.01])
    controller = LQRController(A, B, Q, R)
    
    # 初始化车辆状态为轨迹起点
    initial_state = trajectory[0].copy()
    current_state = initial_state.copy()
    vehicle.reset(initial_state)
    
    # 记录状态历史
    states_history = [current_state.copy()]
    controls_history = []
    
    # 运行仿真
    print("\n开始仿真...")
    for i, t in enumerate(time_points[:-1]):
        # 获取参考状态
        reference_state = trajectory[i].copy()
        
        # 计算控制输入
        u = controller.compute_control_input(current_state, reference_state)
        controls_history.append(u.copy())
        
        # 更新车辆状态
        current_state = vehicle.step(u).copy()
        states_history.append(current_state.copy())
        
        # 打印状态（仅前几步和每50步）
        if i < 5 or i % 50 == 0:
            print(f"Step {i} (t={t:.1f}s):")
            print(f"  参考状态: {reference_state}")
            print(f"  当前状态: {current_state}")
            print(f"  控制输入: {u}")
            pos_error = np.linalg.norm(reference_state[:2] - current_state[:2])
            print(f"  位置误差: {pos_error:.4f}")
    
    # 转换为numpy数组
    states_history = np.array(states_history)
    controls_history = np.array(controls_history)
    
    # 计算误差
    position_errors = trajectory[:len(states_history), :2] - states_history[:, :2]
    distance_errors = np.sqrt(np.sum(position_errors**2, axis=1))
    rmse = np.sqrt(np.mean(distance_errors**2))
    max_error = np.max(distance_errors)
    
    print("\n仿真完成!")
    print(f"RMSE: {rmse:.4f}")
    print(f"最大误差: {max_error:.4f}")
    
    # 绘制结果
    plt.figure(figsize=(12, 10))
    
    # 绘制轨迹
    plt.subplot(2, 2, 1)
    plt.plot(trajectory[:, 0], trajectory[:, 1], 'b-', label='参考轨迹')
    plt.plot(states_history[:, 0], states_history[:, 1], 'r--', label='实际轨迹')
    plt.xlabel('X位置')
    plt.ylabel('Y位置')
    plt.title('轨迹跟踪')
    plt.legend()
    plt.grid(True)
    plt.axis('equal')
    
    # 绘制位置随时间变化
    plt.subplot(2, 2, 2)
    plt.plot(time_points[:len(states_history)], trajectory[:len(states_history), 0], 'b-', label='X参考')
    plt.plot(time_points[:len(states_history)], states_history[:, 0], 'r--', label='X实际')
    plt.plot(time_points[:len(states_history)], trajectory[:len(states_history), 1], 'g-', label='Y参考')
    plt.plot(time_points[:len(states_history)], states_history[:, 1], 'm--', label='Y实际')
    plt.xlabel('时间(s)')
    plt.ylabel('位置')
    plt.title('位置随时间变化')
    plt.legend()
    plt.grid(True)
    
    # 绘制误差
    plt.subplot(2, 2, 3)
    plt.plot(time_points[:len(distance_errors)], distance_errors)
    plt.xlabel('时间(s)')
    plt.ylabel('误差')
    plt.title(f'跟踪误差 (RMSE: {rmse:.4f}, 最大误差: {max_error:.4f})')
    plt.grid(True)
    
    # 绘制控制输入
    plt.subplot(2, 2, 4)
    plt.plot(time_points[:len(controls_history)], controls_history[:, 0], 'r-', label='ax')
    plt.plot(time_points[:len(controls_history)], controls_history[:, 1], 'b-', label='ay')
    plt.xlabel('时间(s)')
    plt.ylabel('控制输入')
    plt.title('控制输入随时间变化')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    
    # 创建保存目录
    if not os.path.exists('results'):
        os.makedirs('results')
    
    plt.savefig('results/fixed_tracking_test.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return states_history, trajectory, controls_history, rmse

if __name__ == "__main__":
    # 设置绘图样式
    set_plot_style()
    
    # 运行测试
    states_history, trajectory, controls_history, rmse = test_simple_circle()
    
    print("\n" + "=" * 50)
    print("测试完成！")
    print(f"轨迹跟踪RMSE: {rmse:.4f}")
    print("=" * 50)
    
    # 如果轨迹跟踪效果良好，给出运行主程序的建议
    if rmse < 1.0:
        print("\n轨迹跟踪效果良好！")
        print("您可以使用以下命令运行完整程序：")
        print("python main.py --q-pos 1000 --q-vel 100 --r-acc 0.01")
    else:
        print("\n轨迹跟踪仍有误差，建议进一步调整参数")
        print("请检查终端输出，查找可能的问题") 